package pub.cleangao.mybatis.springboot;

import java.util.HashMap;
import java.util.Map;


public class DataSourceContext {


    // 旧表名与新表名对应; key: oldTableName, value: newTableName
    private static final ThreadLocal<Map<String, String>> tableCurrentThread = new InheritableThreadLocal<>();
    // 是否使用新表名
    private static final ThreadLocal<Boolean> tableNameChangeThread = new InheritableThreadLocal<>();
    // Datasource key
    private static final ThreadLocal<String> dataSourceKeyThread = new InheritableThreadLocal<>();

    /**
     * 获取分库分表新表名
     *
     * @param tableName 当前表名
     * @return 返回新表名
     */
    public static String getNewTableName(String tableName) {
        Map<String, String> map = tableCurrentThread.get();
        if (map == null || map.isEmpty()) {
            // 分库分表的当前线程存储为空，
            // 直接返回原表名
            return tableName;
        }
        String newName = map.get(tableName);
        if (newName == null) {
            return tableName;
        }
        return newName;
    }

    /**
     * 设置旧表名与新表名映射
     *
     * @param tableName    旧表名
     * @param newTableName 新表名
     */
    public static void setTableName(String tableName, String newTableName) {
        if (!tableName.equals(newTableName)) {
            tableNameChangeThread.set(true);
            Map<String, String> map = tableCurrentThread.get();
            if (map == null) {
                map = new HashMap<>();
            }
            map.put(tableName, newTableName);
        }
    }


    /**
     * 是否启用分库分表
     */
    protected static boolean isTableNameChange() {
        return tableNameChangeThread.get() != null && tableNameChangeThread.get();
    }


    /**
     * 当前数据源标识
     */
    public static String getDataSourceKey() {
        String current = dataSourceKeyThread.get();
        if (current == null || current.isEmpty()) {
            return CleavageConst.defaultDataSourceKey;
        }
        return current;
    }

    public static void setDataSourceKey(String dataSourceKey) {
        dataSourceKeyThread.set(dataSourceKey);
    }

    public static void clearDataSourceKey() {
        dataSourceKeyThread.remove();
    }

    public static void clearNewTableName() {
        tableCurrentThread.remove();
        tableNameChangeThread.remove();
    }

}
